
#include "include/h264_dec.hpp"
#include "rclcpp/rclcpp.hpp"
// enum AVPixelFormat get_format_cuda(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts)
// {
// 	const enum AVPixelFormat *p;

// 	for (p = pix_fmts; *p != -1; p++)
// 	{
// 		if (*p == AV_PIX_FMT_CUDA) return *p;
// 	}

// 	fprintf(stderr, "Failed to get HW surface format.\n");
// 	return AV_PIX_FMT_NONE;
// }

h264_dec::h264_dec()
{
}

h264_dec::~h264_dec()
{
	if (this->prepared) close_stream();
}

/**
 * @brief close rtsp stream and free all data structs
 *
 * @return int always 0
 */
int h264_dec::close_stream()
{
	this->prepared = 0;
	sws_freeContext(sws_ctx);
	av_frame_unref(this->rgb_frame);
	av_packet_free(&this->packet);
	avcodec_free_context(&this->av_codec_ctx);
	avformat_close_input(&if_ctx);
	return 0;
}

/**
 * @brief open rtsp stream and prepare a decoder for the stream
 *
 * @param input_url rtsp stream url
 *
 * @return int 0 on success, <0 if failed
 */
int h264_dec::open_stream(std::string input_url, std::string decoder, std::string scale_size)
{
	int ret;
	std::cout << "trying to open input with tcp" << std::endl;
	AVDictionary *options = NULL;
	av_dict_set(&options, "rtsp_transport", "tcp", 0);
	av_dict_set(&options, "probesize", "512", 0);
	//av_dict_set(&options, "buffer_size", "102400", 0);
	//av_dict_set(&options, "stimeout", "2000000", 0);
	//av_dict_set(&options, "max_delay", "500000", 0);
	// open input stream
	ret = avformat_open_input(&this->if_ctx, input_url.c_str(), NULL, &options);
	if (ret < 0)
	{
		std::cerr << "failed to open rtsp url: " << ret << std::endl;
		return -1;
	}
	std::cout << "open input stream" << std::endl;

	// get stream info
	std::cout << "finding stream information, this may take a while..." << std::endl;
	ret = avformat_find_stream_info(this->if_ctx, NULL);
	if (ret < 0)
	{
		std::cerr << "failed to find stream info: " << ret << std::endl;
		return -1;
	}
	std::cout << "found stream info" << std::endl;

	// find a video stream
	ret = av_find_best_stream(this->if_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0);
	if (ret < 0)
	{
		std::cerr << "failed to find a video stream from source: " << ret << std::endl;
		return -1;
	}
	this->vstream_no = ret;
	std::cout << "found video stream" << std::endl;

	// for jetson, can only use the h264_nvv4l2dec decoder for hardware decoding
	this->av_codec = (AVCodec *)avcodec_find_decoder_by_name(decoder.c_str());
	if (this->av_codec == NULL)
	{
		fprintf(stderr,
				"failed to find decoder \"%s\", run \"ffmpeg -decoders | grep "
				"h264\" to see all supported h264 decoders",
				decoder.c_str());
		return -1;
	}

	// alloc codec context
	this->av_codec_ctx = avcodec_alloc_context3(this->av_codec);
	if (this->av_codec_ctx == NULL)
	{
		std::cerr << "failed to alloc codec context" << std::endl;
		return -1;
	}
	std::cout << "alloc codec context" << std::endl;

	// pass stream info
	ret = avcodec_parameters_to_context(this->av_codec_ctx,
										this->if_ctx->streams[this->vstream_no]->codecpar);
	if (ret < 0)
	{
		std::cerr << "failed to pass parameters" << std::endl;
		return -1;
	}
	std::cout << "passed parameters to decoder" << std::endl;
	// // add hardware device contex
	// this->av_codec_ctx->get_format = get_format_cuda;
	// ret = av_hwdevice_ctx_create(&this->av_hwdev_ctx, AV_HWDEVICE_TYPE_CUDA,
	// 							 render_path.c_str(), NULL, 0);
	// if (ret < 0)
	// {
	// 	std::cerr << "failed to create hwdevice ctx" << std::endl;
	// 	return -1;
	// }
	// this->av_codec_ctx->hw_device_ctx = av_buffer_ref(this->av_hwdev_ctx);
	// std::cout << "found hardware device" << std::endl;
	// open decoder
	ret = avcodec_open2(this->av_codec_ctx, this->av_codec, NULL);
	if (ret < 0)
	{
		std::cerr << "failed to open decoder: " << ret << std::endl;
		return -1;
	}
	std::cout << "open decoder" << std::endl;

	this->packet = av_packet_alloc();
	// filters

	int scale_w = 320, scale_h = 320;
	sscanf(scale_size.c_str(), "%dx%d", &scale_w, &scale_h);

	this->rgb_frame = av_frame_alloc();
	if (this->rgb_frame == NULL)
	{
		std::cerr << "failed to alloc rgb frame" << std::endl;
		return -1;
	}

	this->av_frame = av_frame_alloc();
	if (this->av_frame == NULL)
	{
		std::cerr << "failed to alloc av frame" << std::endl;
		return -1;
	}

	ret = av_image_alloc(this->rgb_frame->data, rgb_frame->linesize, scale_w, scale_h,
						 AV_PIX_FMT_BGR24, 1);
	if (ret < 0)
	{
		std::cerr << "failed to fill rgb frame with BGR24 " << scale_size << std::endl;
		return -1;
	}
	this->rgb_frame->width = scale_w;
	this->rgb_frame->height = scale_h;
	this->rgb_frame->format = AV_PIX_FMT_BGR24;
	
	this->prepared = 1;
	std::cout << "decoder prepared now" << std::endl;
	return 0;
}

int h264_dec::decode_frame(uint8_t *data, int len, int *width, int *height, int *step)
{
	int ret;
	auto t = rclcpp::Clock().now();
	static int frame_cnt;
	static double cnt_start;
	static bool wait_key_frame = true;

	while (true)
	{
		av_frame_unref(this->av_frame);
		ret = avcodec_receive_frame(this->av_codec_ctx, this->av_frame);
		if (ret == 0)
		{
			// got frame from decoder, try scale and output
			frame_cnt++;
			do
			{
				t = rclcpp::Clock().now();
				printf("decoded frame: %dx%d, pkt_dts %ld, pts %ld, fps %.3lf, time %lf\n",
					   this->av_frame->width, this->av_frame->height, this->av_frame->pkt_dts,
					   this->av_frame->pts, frame_cnt / (t.seconds() - cnt_start), t.seconds());
				if (t.seconds() - cnt_start >= 1)
				{
					frame_cnt = 0;
					cnt_start = t.seconds();
				}
				this->sws_ctx = sws_getContext(
					this->av_codec_ctx->width, this->av_codec_ctx->height, (AVPixelFormat)this->av_frame->format,
					this->rgb_frame->width, this->rgb_frame->height, AV_PIX_FMT_BGR24,
					SWS_BICUBIC, nullptr, nullptr, nullptr);
				if (this->sws_ctx == NULL)
				{
					fprintf(stderr, "failed to get sws context: %dx%d, %d -> %dx%d, %d",
							this->av_codec_ctx->width, this->av_codec_ctx->height,
							this->av_codec_ctx->pix_fmt, this->rgb_frame->width,
							this->rgb_frame->height, AV_PIX_FMT_BGR24);
					return -1;
				}

				sws_scale(this->sws_ctx, (uint8_t const *const *)this->av_frame->data,
						  this->av_frame->linesize, 0, this->av_frame->height,
						  this->rgb_frame->data, this->rgb_frame->linesize);
				t = rclcpp::Clock().now();
				printf("scaled: %dx%d, output: size %d, time %lf\n", this->rgb_frame->width,
					   this->rgb_frame->height,
					   this->rgb_frame->height * this->rgb_frame->linesize[0], t.seconds());

				// output
				if (this->rgb_frame->linesize[0] * this->rgb_frame->height > len)
				{
					std::cerr << "image size too large: linesize " << this->rgb_frame->linesize
							  << ", height " << this->rgb_frame->height << std::endl;
					break;
				}
				*height = this->rgb_frame->height;
				*width = this->rgb_frame->width;
				*step = this->rgb_frame->linesize[0];
				memcpy(data, this->rgb_frame->data[0], *height * *step);
				ret = 0;
			} while (0);
			return ret;
		}
		else if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
		{
			av_packet_unref(this->packet);

			ret = av_read_frame(this->if_ctx, this->packet);
			t = rclcpp::Clock().now();
			if (ret < 0 && ret != AVERROR(EAGAIN))
			{
				printf("failed to read frame from input (%d), time %lf\n", ret, t.seconds());
				return -1;
			}
			printf("read packet from stream, time %lf\n", t.seconds());
			// if (this->packet->stream_index != this->vstream_no || this->packet->flags != 1)
			if (this->packet->flags == 1)
			{
				wait_key_frame = false;
			}
			if (wait_key_frame || this->packet->stream_index != this->vstream_no)
			{
				// not video frame
				continue;
			}
			
			// send to decoder
			ret = avcodec_send_packet(this->av_codec_ctx, this->packet);
			t = rclcpp::Clock().now();
			if (ret < 0)
			{
				printf("failed to send packet to decoder (%d), flag %d, dts %ld, pts %ld, time %lf\n",
				 ret, this->packet->flags, this->packet->dts, this->packet->pts, t.seconds());
				return -1;
			}
			printf("sent packet to decoder, flag %d, dts %ld, pts %ld, time %lf\n",
				 this->packet->flags, this->packet->dts, this->packet->pts, t.seconds());
		}
		else
		{
			// eof or other
			printf("got unexpected return value %d\n", ret);
		}
	}
	return -1;
}
